package com.cat.dataStructure02;

import java.util.PriorityQueue;

/**
 * @author 曲大人的喵
 * @description https://leetcode.cn/problems/choose-k-elements-with-maximum-sum/description/
 * @create 2025/10/6 19:47
 * @since JDK17
 */

public class Solution02 {
    public long[] findMaxSum(int[] nums1, int[] nums2, int k) {
        int n = nums1.length, last = 0;
        long sum = 0;
        long[] ans = new long[n];
        PriorityQueue<int[]> heap1 = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        PriorityQueue<Integer> heap2 = new PriorityQueue<>();
        for (int i = 0; i < n; i++) {
            heap1.add(new int[]{nums1[i], i});
        }
        while (!heap1.isEmpty()) {
            var x = heap1.poll();
            int idx = x[1], next = heap1.isEmpty() ? 0 : heap1.peek()[0];;
            ans[idx] = sum;
            last = Math.max(last, nums2[idx]);
            if (next <= x[0]) {
                continue;
            }
            sum += last;
            heap2.add(last);
            if (heap2.size() > k) {
                sum -= heap2.poll();
            }
            last = 0;
        }

        return ans;
    }
}
